import numpy as np
import os
import cv2
import pandas as pd
os.environ["KERAS_BACKEND"] = "torch"  # 使用torch 作为后端
from keras.models import load_model, Model
from keras.layers import (
    Dense,
    Input,
    Conv2D,
    Concatenate,
    AveragePooling2D,
    GlobalAveragePooling2D,
)
from keras.losses import MeanSquaredError
from sklearn.model_selection import train_test_split
from keras.callbacks import ModelCheckpoint
from utils.load_data import Pose
from sklearn.metrics import r2_score

# 加载数据
pose_data_list = Pose.load_data("C:/Users/kang_/Desktop/data")
X_RULER = np.array(
    [i.body_size(128 / i.img_width) for i in pose_data_list if i.weight > 0],
    dtype=np.float32,
)
X_DEPTH = np.array(
    [
        cv2.cvtColor(i.depth_img(128 / i.img_width), cv2.COLOR_GRAY2BGR)
        for i in pose_data_list
        if i.weight > 0
    ]
)
Y = np.array([i.weight for i in pose_data_list if i.weight > 0], dtype=np.float32)
X_ruler_train, X_ruler_val, X_depth_train, X_depth_val, Y_train, Y_val = (
    train_test_split(X_RULER, X_DEPTH, Y, test_size=0.1)
)

# 定义模型
img_shape = X_DEPTH[0].shape
input_image = Input(shape=img_shape, name="image_input")
input_array = Input(shape=(X_RULER[0].size,), name="ruler_input")

# 深度输入到卷积层
img = Conv2D(16, (3, 3), activation="relu")(input_image)
img = AveragePooling2D(2, 2)(img)
img = Conv2D(32, (3, 3), activation="relu")(img)
img = AveragePooling2D(2, 2)(img)
img = Conv2D(64, (3, 3), activation="relu")(img)
img = GlobalAveragePooling2D()(img)
img = Dense(1024, activation="relu")(img)
# 体尺到隐藏层
ruler = Dense(1024, activation="relu")(input_array)
# 合并两个输入
combined = Concatenate()([img, ruler])
hidden = Dense(2048, activation="relu")(combined)
output = Dense(1)(hidden)

# 创建模型
model = Model(inputs=[input_image, input_array], outputs=output)

# 编译模型
model.compile(optimizer="adam", loss=MeanSquaredError(), metrics=["accuracy"])
model.summary()

checkpoint = ModelCheckpoint(
    "best_model_v2.keras", monitor="loss", save_best_only=True, mode="min"
)

# 模型训练
activations = model.fit(
    [X_depth_train, X_ruler_train],
    Y_train,
    epochs=200,
    batch_size=64,
    verbose=2,
    validation_split=0.2,
    callbacks=[checkpoint],
)

model_best = load_model("./best_model_v2.keras")
pred = model_best.predict([X_depth_val, X_ruler_val])
print(r2_score(Y_val, pred))

v = np.array(Y_val).flatten()
p = np.array(pred).flatten()
df = pd.DataFrame({ 'val': v, 'pred': p })
df.to_csv('output.csv', index=False)